__author__ = 'carlxie'

import numpy as np
from LayerType import LayerType

# remember our old friend ??
def sigmoid(s):
    return 1.0 / (1.0 + np.exp(-s))

def uniform_rand(r, shape):
    return 2 * r * np.random.random(shape) - r

## tanh is just a scaled and shifted sigmoid function
def tanh(s):
    return 2 * sigmoid(2 * s) - 1

def der_tanh(s):
    t = tanh(s)
    return 1 - t ** 2

def leru(x):
    return max(0,x)

def leru_helper(x):
    if x > 0:return 1
    else:return 0


def vec_output(y,size):
    out = np.ones((size,1))
    out *= -1
    if size > 1:
        out[y,0] = 1
    else:
        if y > 0:
            out[0,0] = 1
    return out

# vectorize our util function
vec_tanh = np.vectorize(tanh)
vec_der_tanh = np.vectorize(der_tanh)
vec_leru = np.vectorize(leru)

def get_weights(nn):
    w = []
    for layer in nn.layers:
        if layer.layerType != LayerType.MAX_POOL_TYPE and layer.layerType != LayerType.INPUT_TYPE:
            w.append(layer.weights)
    return flat_list(w)

def get_biases(nn):
    biases = []
    for layer in nn.layers:
        if layer.layerType != LayerType.MAX_POOL_TYPE and layer.layerType != LayerType.INPUT_TYPE:
            biases.append(layer.biases)
    return flat_list(biases)

def get_flat_size(shape):
    return reduce(lambda x,y:x*y,shape,1)

def flat_list(ls):
    fls = np.array([])
    for i in range(len(ls)):
        flat_size = get_flat_size(ls[i].shape)
        fls = np.append(fls, ls[i].reshape(flat_size))
    return fls

##
#  shapes is something like [(2,3),(3,2,3),(1,2,3)...]
##
def reconstruct(ls, shapes):
    ptr = 0
    wls = []
    bls = []
    for i in range(len(shapes)):
        flat_size = get_flat_size(shapes[i])
        end_ptr = ptr + flat_size
        wls.append(ls[ptr:end_ptr].reshape(shapes[i]))
        ptr += flat_size
    for i in range(len(shapes)):
        flat_size = shapes[i][-1]
        end_ptr = ptr + flat_size
        if len(shapes[i]) == 4:
            bls.append(ls[ptr:end_ptr].reshape((1,1,flat_size)))
        elif len(shapes[i]) == 2:
            bls.append(ls[ptr:end_ptr].reshape((flat_size,1)))
        ptr += flat_size
    return (wls,bls)

def get_shapes(nn):
    shapes = []
    for layer in nn.layers:
        if layer.layerType != LayerType.MAX_POOL_TYPE and layer.layerType != LayerType.INPUT_TYPE:
            shapes.append(layer.weights.shape)
    return shapes

def compute_num_grads(nn,x,y,shapes,weights):
    e = 1e-4
    perturb = np.zeros(weights.shape)
    num_grad = np.zeros(weights.shape)
    for i in range(len(weights)):
        perturb[i] = e
        w1,b1 = reconstruct(weights + perturb, shapes)
        w2,b2 = reconstruct(weights - perturb, shapes)
        perturb[i] = 0
        loss1 = nn.cal_loss(x,y,w1,b1)
        loss2= nn.cal_loss(x,y,w2,b2)
        num_grad[i] = (loss1 - loss2) / (2 * e)
    return num_grad